package org.dromara.common.web.filter

import cn.hutool.core.io.IoUtil
import cn.hutool.core.util.StrUtil
import cn.hutool.http.HtmlUtil
import jakarta.servlet.ReadListener
import jakarta.servlet.ServletInputStream
import jakarta.servlet.http.HttpServletRequest
import jakarta.servlet.http.HttpServletRequestWrapper
import org.dromara.common.core.utils.StringUtils
import org.springframework.http.HttpHeaders
import org.springframework.http.MediaType
import java.io.IOException
import java.nio.charset.StandardCharsets

/**
 * XSS过滤处理
 *
 * @author LikeYouDo
 */
class XssHttpServletRequestWrapper(request: HttpServletRequest) : HttpServletRequestWrapper(request) {
    override fun getParameter(name: String): String? {
        val value = super.getParameter(name) ?: return null
        return HtmlUtil.cleanHtmlTag(value).trim()
    }

    override fun getParameterMap(): MutableMap<String, Array<String?>?> {
        val valueMap = super.getParameterMap()
        if (valueMap.isNullOrEmpty()) {
            return valueMap
        }
        // 避免某些容器不允许改参数的情况 copy一份重新改
        val map: MutableMap<String, Array<String?>?> = HashMap(valueMap.size)
        map.putAll(valueMap)
        for (entry in map.entries) {
            val values = entry.value
            if (values != null) {
                val length = values.size
                val escapseValues = arrayOfNulls<String>(length)
                for (i in 0..<length) {
                    // 防xss攻击和过滤前后空格
                    escapseValues[i] = HtmlUtil.cleanHtmlTag(values[i]).trim()
                }
                map[entry.key] = escapseValues
            }
        }
        return map
    }

    override fun getParameterValues(name: String): Array<String?> {
        val values = super.getParameterValues(name)
        if (values.isNullOrEmpty()) {
            return values
        }
        val length = values.size
        val escapseValues = arrayOfNulls<String>(length)
        for (i in 0..<length) {
            // 防xss攻击和过滤前后空格
            escapseValues[i] = HtmlUtil.cleanHtmlTag(values[i]).trim()
        }
        return escapseValues
    }

    @Throws(IOException::class)
    override fun getInputStream(): ServletInputStream {
        // 非json类型，直接返回
        if (!isJsonRequest()) {
            return super.getInputStream()
        }

        // 为空，直接返回
        var json = StrUtil.str(IoUtil.readBytes(super.getInputStream(), false), StandardCharsets.UTF_8)
        if (json.isNullOrEmpty()) {
            return super.getInputStream()
        }

        // xss过滤
        json = HtmlUtil.cleanHtmlTag(json).trim()
        val jsonBytes = json.toByteArray(StandardCharsets.UTF_8)
        val bis = IoUtil.toStream(jsonBytes)
        return object : ServletInputStream() {
            override fun isFinished(): Boolean {
                return true
            }

            override fun isReady(): Boolean {
                return true
            }

            @Throws(IOException::class)
            override fun available(): Int {
                return jsonBytes.size
            }

            override fun setReadListener(readListener: ReadListener?) {
            }

            @Throws(IOException::class)
            override fun read(): Int {
                return bis.read()
            }
        }
    }

    /**
     * 是否是Json请求
     */
    fun isJsonRequest(): Boolean {
        val header = super.getHeader(HttpHeaders.CONTENT_TYPE)
        return StringUtils.startsWithIgnoreCase(header, MediaType.APPLICATION_JSON_VALUE)
    }
}
